# coding=utf-8

from PyPDF2 import PdfFileWriter, PdfFileReader
from copy import copy




def create_watermark(input_pdf, output, watermark):
    '''
    添加水印
    :param input_pdf: 原始pdf
    :param output: 输出的pdf
    :param watermark: 水印
    :return:
    '''
    # try:
    watermark_obj = PdfFileReader(watermark)
    watermark_page = watermark_obj.getPage(0)

    pdf_reader = PdfFileReader(input_pdf)
    pdf_writer = PdfFileWriter()

    # TODO： 给所有页面添加水印
    for page in range(pdf_reader.getNumPages()):
        page = pdf_reader.getPage(page)
        new_watermark_page = copy(watermark_page)
        new_watermark_page.mergePage(page)
        pdf_writer.addPage(new_watermark_page)

    with open(output, 'wb') as out:
    # with open(output, 'x') as out:
        pdf_writer.write(out)

# except Exception as e:
#     print("e",e )
#     pass


if __name__ == '__main__':
    # create_watermark(input_pdf='09宁乡市排水管网地理信息系统项目培训方案.pdf',output='watermarked.pdf',watermark='中南院水印.pdf')
    input_pdf = '宁乡市排水管网地理信息系统项目培训方案.pdf'
    output = 'watermarked.pdf'
    watermark = '中南院水印.pdf'
    create_watermark(input_pdf, output, watermark)
